# import asyncio
# import websockets
# import sys
#
# IP = '127.0.0.1'
# if len(sys.argv) > 1:
#     IP = str(sys.argv[1])
#     print(IP)
#
# PORT = '1234'
# if len(sys.argv) > 2:
#     PORT = str(sys.argv[2])
#     print(PORT)
#
# URI = f"ws://{IP}:{PORT}"
#
# # 异步操作，定义一个异步函数进行消息接收和发送
# async def conn():
#     uri = URI  # WebSocket 服务器的 URL
#
#     async with websockets.connect(
#         uri,ping_interval=None
#     ) as websocket:
#         while True:
#             msg = input(f"{IP}:{PORT} > ")
#             if msg and not msg.isspace():
#                 if msg.lower() in ("exit", "quit"):
#                     await websocket.send(msg)  # 发送消息
#                     print("Exiting Connecting!")
#                     break  # 退出连接循环
#                 # 发送消息
#                 await websocket.send(msg)
#                 # 接收消息
#                 response = await websocket.recv()
#                 print(response)
#
#
#
#
# # 获取默认的事件循环，并运行连接
# loop = asyncio.get_event_loop()
# loop.run_until_complete(conn())
#

import asyncio
import websockets
import sys
import keyboard

async def subclient(websocket, msg, nowlen):
    while True:
        try:
            await asyncio.sleep(1)
            await websocket.send(msg)
            response = await websocket.recv()
            if "(error)" in response:
                print(response)
                break;
            else:
                ls = response.split(" ")
                tmp = []
                if nowlen != len(ls):
                    for i in range(nowlen, len(ls)):
                        res = f"{i}){ls[i]}"
                        tmp.append(res)
                    nowlen = len(ls)
                    tmp = "\n".join(tmp)
                    print(tmp)
        except KeyboardInterrupt:
            print("ctrl+x to exit")
        if keyboard.is_pressed('ctrl') and keyboard.is_pressed('x'):
            print("exiting...")
            break

# 标识连接
class Connect:
    def __init__(self,ip='127.0.0.1',port='1234'):
        self.ip = ip
        self.port = port
        self.uri = f"ws://{self.ip}:{self.port}"
        self.flag = False
        self.msg = ''
        self.rdbId = '0'

    def setURI(self,ip,port):
        self.ip = ip
        self.port = port
        self.uri = f"ws://{self.ip}:{self.port}"
        return self.uri

    def getURI(self):
        return self.uri


"""这里在启动时指定对应的ip和port"""
# connect = Connect('112.124.19.9')  # 阿里云
connect = Connect()

if len(sys.argv) > 1:
    connect.ip = str(sys.argv[1])

PORT = '1234'
if len(sys.argv) > 2:
    connect.port = str(sys.argv[2])


async def conn():
    while True:
        # 创建 WebSocket 连接
        async with websockets.connect(
            connect.getURI(), ping_interval=None
        ) as websocket:
            while True:

                if connect.flag:  # 如果切换过服务器数据库，需要重新执行
                    connect.flag = False
                else:
                    if connect.rdbId == '0':
                        connect.msg = input(f"{connect.ip}:{connect.port}> ")
                    else:
                        connect.msg = input(f"{connect.ip}:{connect.port}({connect.rdbId})> ")
                msg = connect.msg
                if msg and not msg.isspace():
                    if msg.lower() in ("exit", "quit"):
                        await websocket.send(msg)
                        print("Exiting Connecting!")
                        return

                    cmd_list = msg.split(" ")
                    if cmd_list[0] == "subscribe":
                        nowlen = 0
                        print("reading...")
                        await subclient(websocket, msg, nowlen)
                    elif 'select' == cmd_list[0]:
                        await websocket.send(msg)
                        # 接收消息
                        response = await websocket.recv()
                        if response in "OK":
                            rdb_id = msg.split()[1]
                            connect.rdbId = rdb_id
                        print(response)
                    else:
                        await websocket.send(msg)
                        # 接收消息
                        response = await websocket.recv()
                        print(response)
                        # 判断服务器是否返回新的地址
                        if "REDIRECT" in response:
                            addr = response.split()[-1]
                            new_ip, new_port = addr.split(':')[:]
                            connect.setURI(new_ip, new_port)
                            print(f"Switching to new address: {connect.ip}:{connect.port}")
                            print(f"{connect.ip}:{connect.port}> {msg}")
                            connect.flag = True
                            # 关闭当前连接，重新连接新的地址
                            break


# 使用默认的事件循环运行连接
asyncio.run(conn())
